#include <algorithm/analytical_model/algorithm/cuda/Utils.cuh>
#include <algorithm/analytical_model/algorithm/cuda/Utils.h>
#include <algorithm/analytical_model/algorithm/pipeline/Utils.h>
#include <utils/Format.h>
#include <utils/Utils.h>

#include <cuda_runtime_api.h>
#include <thrust/async/for_each.h>
#include <thrust/async/reduce.h>
#include <thrust/device_ptr.h>
#include <thrust/device_vector.h>
#include <thrust/extrema.h>
#include <thrust/iterator/counting_iterator.h>

using namespace solar;
using namespace solar::cuda::inter_op;

// namespace solar::cuda::internal
// {
//     __global__ static void sumFluxMap(int2 resolution, float *&result, const float *input)
//     {
//         int tx = threadIdx.x;
//         int ty = threadIdx.y;
//         int bw = blockDim.x;
//         int bh = blockDim.y;
//         int x = blockIdx.x * bw + tx;
//         int y = blockIdx.y * bh + ty;
//         if (x >= resolution.x || y >= resolution.y)
//         {
//             return;
//         }
//         result[y * resolution.x + x] += input[y * resolution.x + x];
//     }
// }

void solar::cuda::sumFluxMap(const Int2& resolution, float*& result, float* input,
                             cudaStream_t stream, bool delete_input)
{
    const int total_resolution = resolution.x * resolution.y;
    thrust::async::for_each(thrust::device.on(stream), thrust::make_counting_iterator(0),
                            thrust::make_counting_iterator(total_resolution),
                            [input, result] __device__(const int& index)
                            { result[index] += input[index]; });
}

void solar::cuda::divFluxMap(const solar::cuda::inter_op::Int2& resolution, float*& result,
                             const float& num, cudaStream_t stream)
{
    const int total_resolution = resolution.x * resolution.y;
    thrust::async::for_each(thrust::device.on(stream), thrust::make_counting_iterator(0),
                            thrust::make_counting_iterator(total_resolution),
                            [num, result] __device__(const int& index)
                            { result[index] = result[index] / num; });
}

auto solar::findFluxMapMax(const Int2& resolution, float* input) -> std::tuple<float, std::uint64_t>
{
    const thrust::device_ptr<float> input_ptr(input);
    auto pixel_max_iter = thrust::max_element(input_ptr, input_ptr + resolution.x * resolution.y);
    auto index = thrust::distance(input_ptr, pixel_max_iter);
    return {*pixel_max_iter, index};
}

auto solar::totalEnergyOnFluxMap(const Int2& resolution, float* input) -> float
{
    const thrust::device_ptr<float> input_ptr(input);
    auto pixel_sum = thrust::reduce(input_ptr, input_ptr + size_t(resolution.x * resolution.y),
                                    0.0F, thrust::plus<float>());
    return pixel_sum;
}
